#include <Eigen/Dense>
#include <chrono>
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <string>

#include <CLI/App.hpp>
#include <CLI/CLI.hpp>

#include "IVPSolverFactory.hpp"

constexpr double mu = 0.012277471;

Eigen::Vector<double, 6> f(Eigen::Vector<double, 6> u, double t) {
  Eigen::Vector<double, 6> du;
  du(0) = u(3);
  du(1) = u(4);
  du(2) = u(5);
  du(3) = 2.0 * u(4) + u(0) -
          mu * (u(0) + mu - 1) /
              pow(pow(u(1), 2) + pow(u(2), 2) + pow(u(0) + mu - 1, 2), 1.5) -
          (1 - mu) * (u(0) + mu) /
              pow(pow(u(1), 2) + pow(u(2), 2) + pow(u(0) + mu, 2), 1.5);
  du(4) = -2.0 * u(3) + u(1) -
          mu * u(1) /
              pow(pow(u(1), 2) + pow(u(2), 2) + pow(u(0) + mu - 1, 2), 1.5) -
          (1 - mu) * u(1) /
              pow(pow(u(1), 2) + pow(u(2), 2) + pow(u(0) + mu, 2), 1.5);
  du(5) = -mu * u(2) /
              pow(pow(u(1), 2) + pow(u(2), 2) + pow(u(0) + mu - 1, 2), 1.5) -
          (1 - mu) * u(2) /
              pow(pow(u(1), 2) + pow(u(2), 2) + pow(u(0) + mu, 2), 1.5);
  return du;
};

int main(int argc, char **argv) {
  CLI::App app;
  bool list_solvers = false;
  app.add_flag("-l,--list", list_solvers, "list all usable solvers");
  int type = 1;
  app.add_option("-t,--type", type, "ode type, 1 or 2, default 1");
  double dt = 0.01;
  app.add_option("-s,--step", dt, "step size, default 0.01");
  std::string solverName = "Fehlberg";
  app.add_option("-n,--name", solverName, "solver name, default Fehlberg");
  std::string outputFile = "output.txt";
  app.add_option("-o,--output", outputFile, "output file, default output.txt");
  CLI11_PARSE(app, argc, argv);

  Eigen::Vector<double, 6> u;
  double t0 = 0;
  double t1;

  if (type == 1) {
    u(0) = 0.994;
    u(1) = 0;
    u(2) = 0;
    u(3) = 0;
    u(4) = -2.0015851063790825224;
    u(5) = 0;
    t1 = 17.06521656015796;
  } else if (type == 2) {
    u(0) = 0.879779227778;
    u(1) = 0;
    u(2) = 0;
    u(3) = 0;
    u(4) = -0.379677780949;
    u(5) = 0;
    t1 = 19.140540691377;
  }

  RegisterAllIVPSolvers();
  auto &factory = IVPSolverFactory::getInstance();
  auto solver = factory.createIVPSolver(solverName, 6, dt, t0, t1);
  if (list_solvers) {
    std::cout << "usable solvers:" << std::endl;
    std::cout << "------------------------------" << std::endl;

    auto solvers = factory.getRegisteredSolvers();
    // 排序后显示
    std::sort(solvers.begin(), solvers.end());

    for (const auto &name : solvers) {
      std::cout << name << "\t- " << factory.getSolverDescription(name)
                << std::endl;
    }
    std::cout << "------------------------------" << std::endl;
    std::cout << "All " << solvers.size() << " usable solvers" << std::endl;
    return 0;
  }

  if (!factory.isSolverRegistered(solverName)) {
    std::cerr << "ERR: unknown solver '" << solverName << "'" << std::endl;
    std::cerr << "run witu --list to check all usable solvers" << std::endl;
    return 1;
  }

  std::cout << "Initial condition: " << u.transpose() << std::endl;
  std::cout << "dt: " << dt << std::endl;
  std::cout << "Solver: " << solverName << std::endl;
  std::cout << "Start integral" << std::endl;

  auto start = std::chrono::high_resolution_clock::now();
  solver->solve(u, f);
  auto end = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> elapsed = end - start;

  std::cout << "Finish integral" << std::endl;
  std::cout << "Elapsed time: " << elapsed.count() << " seconds" << std::endl;

  std::ofstream output(outputFile);
  Eigen::MatrixXd U = solver->getSolution();
  for (int i = 0; i < U.rows(); ++i) {
    output << U(i, 0) << " " << U(i, 1) << std::endl;
  }
  std::cout << "Output saved to " << outputFile.c_str() << std::endl;
  return 0;
}